package com.bondar.treemap;

import java.util.Comparator;
import java.util.Iterator;

import com.bondar.interfaces.MyMap;

public class MyTreeMap<Key extends Comparable<Key>> implements MyMap<Key>,
		Comparator<Key> {

	private SimpleEntry<Key> root;
	private int size;
	private Comparator<Key> comparator;

	public MyTreeMap() {
		size = 0;
		comparator = null;
	}

	public MyTreeMap(Comparator<Key> comparator) {
		this.comparator = comparator;
	}

	@Override
	public void clear() {
		root = null;
		size = 0;
	}

	@Override
	public boolean containsKey(Key key) {
		return get(key) == null ? false : true;
	}

	@Override
	public boolean containsValue(Object value) {
		if (value == null) {
			return getValueForNull(root) != null ? true : false;
		}
		return getValue(root, value) != null ? true : false;
	}

	private SimpleEntry<Key> getValueForNull(SimpleEntry<Key> node) {
		if (node == null) {
			return null;
		}

		if (node.getValue() == null) {
			return node;
		}

		SimpleEntry<Key> res;
		if ((res = getValueForNull(node.getLeft())) != null) {
			return res;
		}

		return getValueForNull(node.getRight());
	}

	private Object getValue(SimpleEntry<Key> node, Object value) {
		if (node == null) {
			return null;
		}
		Object result;

		if ((result = node.getValue()) != null) {
			if (result.equals(value)) {
				return result;
			}
		}

		if ((result = getValue(node.getLeft(), value)) != null) {
			return result;
		}

		return getValue(node.getRight(), value);
	}

	@Override
	public Object get(Key key) throws NullPointerException {
		if (key == null) {
			throw new NullPointerException();
		}

		SimpleEntry<Key> x = root;
		while (x != null) {
			int compare = x.getKey().compareTo(key);
			if (compare > 0) {
				x = x.getLeft();
			} else if (compare < 0) {
				x = x.getRight();
			} else if (compare == 0) {
				return x.getValue();
			}
		}
		return null;
	}

	@Override
	public boolean isEmpty() {
		return size == 0;
	}

	@Override
	public Object remove(Key key) {
		if (containsKey(key) == false) {
			throw new NullPointerException();
		}
		
		SimpleEntry<Key> node = root;
		Object value = null;

		if (size == 1) {
			value = root.getValue();
			root = null;
			--size;
			return value;
		}

		node.setRed(true);
		int compare = node.getKey().compareTo(key);

		// find place
		do {
			compare = node.getKey().compareTo(key);

			if (compare > 0) {
				if (node.getLeft() == null
						|| (!isRed(node.getLeft()) && !isRed(node.getLeft()
								.getLeft()))) {
					// paint to red child's
					flipColors(node);
					if (isRed(node.getRight().getLeft())) {
						newRotateRight(node.getRight());
						newRotateLeft(node);
					}
				}
				node = node.getLeft();

			} else {

				if (isRed(node.getLeft())) {
					newRotateRight(node);
				}

				// it is leaf
				if (compare == 0 && node.getRight() == null) {
					value = node.getValue();

					if (node.getParent().getLeft() == node) {
						node.getParent().setLeft(null);
					} else {
						node.getParent().setRight(null);
					}
					
					node = node.getParent();
					break;

				} else if (!isRed(node.getRight())
						&& (node.getRight() == null || !isRed(node.getRight()
								.getLeft()))) {
					flipColors(node);
					if (isRed(node.getLeft().getLeft())) {
						newRotateRight(node);
					}
				}

				if (compare == 0) {
					// find minimal key from right subtree
					SimpleEntry<Key> min = findMin(node.getRight());
					value = node.getValue();

					// swap min node with current
					node.setKey(min.getKey());
					node.setValue(min.getValue());
					// delete right node
					// deleteMin(node.getRight());
					if (node.getRight() == min) {
						node.setRight(null);
					} else {
						min.getParent().setLeft(null);
					}
				} else {
					node = node.getRight();
				}
			}

		} while (compare != 0);

		--size;
		// balance up to root node
		balance(node, null);
		root.setRed(false);
		return value;
	}

	private SimpleEntry<Key> findMin(SimpleEntry<Key> node) {
		if (node.getLeft() == null) {
			return node;
		} else {
			return findMin(node.getLeft());
		}
	}

	@Override
	public int size() {
		return size;
	}

	private void print(SimpleEntry<Key> node) {
		if (node == null) {
			return;
		}

		print(node.getLeft());
		System.out.println(node);
		print(node.getRight());
	}

	public void print() {
		System.out.println("Root: " + root);
		print(root);
	}

	@Override
	public Iterator<Entry<Key>> entryIterator() {

		return new Iterator<Entry<Key>>() {
			SimpleEntry<Key> node = findMin(root);
			SimpleEntry<Key> last;
			int index = size;

			@Override
			public boolean hasNext() {
				return index > 0;
			}

			@Override
			public MyMap.Entry<Key> next() {
				if (node == root && index > 1) {
					node = findMin(root.getRight());
				}
				
				SimpleEntry<Key> entry = node;
				index--;

				if (node.getRight() != null) {
					node = node.getRight();
				} else {
					node = node.getParent();
					if (node == last) {
						node = node.getParent();
					}
				}

				last = entry;
				return entry;
			}

			@Override
			public void remove() {
			}
		};

	}

	private void flipColors(SimpleEntry<Key> node) {
		assert ((!isRed(node)) && isRed(node.getLeft())
				&& isRed(node.getRight()) || (isRed(node))
				&& !isRed(node.getLeft()) && !isRed(node.getRight()));

		node.setRed(!node.isRed());
		node.getLeft().setRed(!node.getLeft().isRed());
		node.getRight().setRed(!node.getRight().isRed());
	}

	private boolean isRed(SimpleEntry<Key> node) {
		if (node == null) {
			return false;
		}
		return node.isRed();
	}

	@Override
	public Object put(Key key, Object value) {

		SimpleEntry<Key> node = root;
		if (node == null) {
			++size;
			root = new SimpleEntry<Key>(key, value, false, null);
			return null;
		}

		SimpleEntry<Key> parent;
		int compare;

		do {
			compare = node.getKey().compareTo(key);
			parent = node;

			if (compare > 0) {
				node = node.getLeft();
			} else if (compare < 0) {
				node = node.getRight();
			} else if (compare == 0) {
				return node.setValue(value);
			}
		} while (node != null);

		// we get to the bottom of the tree
		SimpleEntry<Key> newNode = new SimpleEntry<Key>(key, value, true,
				parent);
		if (compare < 0) {
			parent.setRight(newNode);
		} else {
			parent.setLeft(newNode);
		}

		// balance all nodes from current to root nodes
		balance(parent, null);

		++size;
		// root node should always be black
		root.setRed(false);
		return null;
	}

	private void balance(SimpleEntry<Key> parent, SimpleEntry<Key> till) {

		while (parent != till) {
			if (isRed(parent.getRight()) && !isRed(parent.getLeft())) {
				newRotateLeft(parent);
			}
			if (isRed(parent.getLeft()) && isRed(parent.getLeft().getLeft())) {
				newRotateRight(parent);
			}
			if (isRed(parent.getLeft()) && isRed(parent.getRight()))
				flipColors(parent);

			parent = parent.getParent();

		}

	}

	private void newRotateLeft(SimpleEntry<Key> node) {
		SimpleEntry<Key> x = node.getRight();
		SimpleEntry<Key> childX;

		if (node == root) {
			root = x;
			x.setParent(null);
		} else {
			x.setParent(node.getParent()); 
			// choose right sub-node
			if (node.getParent().getRight() == node) {
				node.getParent().setRight(x);
			} else {
				node.getParent().setLeft(x);
			}
		}

		childX = x.getLeft();
		node.setRight(childX);
		if (childX != null) {
			childX.setParent(node);
		}
		
		x.setLeft(node);
		node.setParent(x);

		x.setRed(node.isRed());
		node.setRed(true);
	}

	private void newRotateRight(SimpleEntry<Key> node) {
		SimpleEntry<Key> x = node.getLeft();
		SimpleEntry<Key> childX;

		if (node == root) {
			root = x;
			x.setParent(null);

		} else {
			x.setParent(node.getParent());
			if (node.getParent().getRight() == node) {
				node.getParent().setRight(x);
			} else {
				node.getParent().setLeft(x);
			}
		}

		childX = x.getRight();
		node.setLeft(childX);
		if (childX != null) {
			childX.setParent(node);
		}
		
		x.setRight(node);
		node.setParent(x);

		x.setRed(node.isRed());
		node.setRed(true);
	}

	@Override
	public int compare(Key o1, Key o2) {
		if (comparator != null) {
			return comparator.compare(o1, o2);
		}
		return o1.compareTo(o2);
	}

	public static class SimpleEntry<Key extends Comparable<Key>> implements
			Entry<Key> {
		private SimpleEntry<Key> left;
		private SimpleEntry<Key> right;
		private SimpleEntry<Key> parent;
		private Object value;
		private Key key;
		private boolean red;

		public SimpleEntry(Key key, Object value, boolean red,
				SimpleEntry<Key> parent) {
			this.key = key;
			this.value = value;
			this.red = red;
			this.parent = parent;
		}

		public boolean isRed() {
			return red;
		}

		public SimpleEntry<Key> getLeft() {
			return left;
		}

		public SimpleEntry<Key> getRight() {
			return right;
		}

		public Object getValue() {
			return value;
		}

		public Object setValue(Object value) {
			Object oldVal = this.value;
			this.value = value;
			return oldVal;
		}

		public Key getKey() {
			return key;
		}

		public void setKey(Key key) {
			this.key = key;
		}

		public void setLeft(SimpleEntry<Key> left) {
			this.left = left;
		}

		public void setRight(SimpleEntry<Key> right) {
			this.right = right;
		}

		public void setRed(boolean red) {
			this.red = red;
		}

		public void setParent(SimpleEntry<Key> node) {
			parent = node;
		}

		public SimpleEntry<Key> getParent() {
			return parent;
		}

		@Override
		public String toString() {
			return String.format("[%s] - %s - %s", key, value,
					red == true ? "RED" : "BLACK");
		}
	}

}
